import pandas as pd
import matplotlib.pyplot as plt
from gluonts.dataset.pandas import PandasDataset
from gluonts.dataset.split import split
from huggingface_hub import hf_hub_download
from uni2ts.eval_util.plot import plot_single
from uni2ts.model.moirai import MoiraiForecast, MoiraiModule
from vnpy.trader.database import get_database
from vnpy_mongodb.mongodb_database import MongodbDatabase
from datetime import datetime, timedelta
# 读取本地数据库的K线数据
database: MongodbDatabase = get_database()
back_price = database.load_back_bar_data(symbol="EURUSD", freq="1H", start=datetime(2023, 3, 1, 0, 0), end=datetime(2023, 9, 30, 0, 0))

# 将数据转换为DataFrame
df = pd.DataFrame(back_price)

# 将datetime列设置为索引
df['datetime'] = pd.to_datetime(df['datetime'])
df.set_index('datetime', inplace=True)

# # 检测 datetime 是否能被 5 分钟整除
# df['is_valid'] = df.index.to_series().apply(lambda x: x.minute % 5 == 0)

# # 筛选出有效的数据
# valid_df = df[df['is_valid']].drop(columns=['is_valid'])

# # 筛选出无效的数据
# invalid_df = df[~df['is_valid']].drop(columns=['is_valid'])

# 确保 DataFrame 的索引具有明确的1小时频率
df = df.asfreq('1H')

# 移除非数值列
valid_df = df.drop(columns=['symbol', 'freq'])


# 转换为 GluonTS 数据集
ds = PandasDataset(dict(valid_df))

# 定义测试集长度
TEST = 100  # 根据你的数据量调整

# 划分训练集和测试集
train, test_template = split(ds, offset=-TEST)

# 定义预测长度和窗口数量
PDT = 20  # 预测长度
windows = TEST // PDT  # 窗口数量
distance = PDT  # 窗口间隔

# 生成滚动窗口评估数据
test_data = test_template.generate_instances(
    prediction_length=PDT,
    windows=windows,
    distance=distance
)

# 定义模型参数
SIZE = "large"  # 模型大小
CTX = 200  # 上下文长度
PSZ = "auto"  # 补丁大小
BSZ = 32  # 批量大小

# 准备预训练模型
model = MoiraiForecast(
    module=MoiraiModule.from_pretrained(f"Salesforce/moirai-1.1-R-{SIZE}"),
    prediction_length=PDT,
    context_length=CTX,
    patch_size=PSZ,
    num_samples=100,
    target_dim=1,
    feat_dynamic_real_dim=ds.num_feat_dynamic_real,
    past_feat_dynamic_real_dim=ds.num_past_feat_dynamic_real,
)

# 创建预测器
predictor = model.create_predictor(batch_size=BSZ)

# 进行预测
forecasts = predictor.predict(test_data.input)

# 获取第一个输入、标签和预测结果
input_it = iter(test_data.input)
label_it = iter(test_data.label)
forecast_it = iter(forecasts)

inp = next(input_it)
label = next(label_it)
forecast = next(forecast_it)

print(forecast)

# 绘制预测结果
plot_single(
    inp,
    label,
    forecast,
    context_length=CTX,
    name="pred",
    show_label=True,
)
plt.show()
